import taichi as ti

ti.init(arch=ti.gpu)

@ti.data_oriented
class SpringSystem:
    def __init__(self,dt):
        self.initialize()
        self.dt = dt
    
    def initialize(self):
        self.x = ti.Vector.field(2,dtype=ti.f32,shape=1)
        self.v = ti.Vector.field(2,dtype=ti.f32,shape=1)
        self.g = ti.Vector([0.0,-10.0])
        self.m = ti.field(dtype=ti.f32,shape=())
        self.m[None] = 1
        self.k = ti.field(dtype=ti.f32,shape=())
        self.k[None] = 200
        self.x[0][0],self.x[0][1] = 0.5,0.8

    @ti.kernel
    def integrate(self):
        # v[t+1] = v[t] + dt* (r*GM/r3)
        self.v[0] += self.dt * (self.g-self.k[None]/self.m[None]*(self.x[0]-ti.Vector([0.5,0.5])))
        # x[t+1] = x[t] + dt*v[t+1]
        self.x[0] += self.dt * self.v[0]
    
    def renderGUI(self):
        self.gui = ti.GUI("spring system",background_color=0x25A6D9)
        while True:
            if self.gui.get_event():
                if self.gui.event.key == self.gui.SPACE and self.gui.event.type == self.gui.PRESS:
                    self.initialize()
            for i in range(100):
                self.integrate()
            self.gui.circle(self.x[0].to_numpy(),radius=20,color=0x8C274C)
            y = self.x[0].to_numpy().tolist()[1]
            x = 0.5
            self.gui.rect([x-0.01+(y-0.5)*0.02,y-0.04],[x+0.01-(y-0.5)*0.02,0],radius=15,color=0x8C28FF)
            self.gui.show()

spring = SpringSystem(dt=0.00005)
spring.renderGUI()
